/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 */

import { Level } from 'log4js';
import path from 'path';
import { Worker } from 'worker_threads';

import { iconv } from '../../../../common/util/iconv/index.js';
import { noop } from '../../../../common/util/noop.js';
import { HvigorLogger } from '../../../log/hvigor-log.js';
import { isHvigorLogWithTime } from '../../../util/time-util';
import { coreParameter } from '../../data/global-core-parameters.js';
import { PoolConstant } from '../constant/constant.js';
import { LogType } from '../enum/log-type.js';
import { TaskState } from '../enum/task-state.js';
import { WorkerState } from '../enum/worker-state.js';
import { WorkLog } from '../log-collection/log.js';
import { ITCB } from '../model/itcb.js';
import { Work } from '../model/work.js';
import { TcbStore } from '../store/tcb-store.js';
import { PoolConfig } from '../worker-pool/pool-config.js';

import { HWorker } from './worker.js';
import { WorkerManager } from './worker-manager.js';
import { ExtendedErrorInfo } from './extended-error-info.js';

/**
 * worker管理器的实现类
 *
 * @since 2022/8/15
 */
export class WorkerManagerImpl implements WorkerManager {
  private readonly maxPoolNum: number;
  private readonly minPoolNum: number;
  private recycleInterval: number;
  private maxIdleTime: number;
  private readonly idleWorkers: Map<number, HWorker>;
  private readonly busyWorkers: Map<number, HWorker>;
  private deadWorkers: Set<number>;
  private timer: NodeJS.Timer;
  private callbacks: Map<string, { fn: Function; input: unknown[] }>;
  private errorCallback: Function = noop;
  private eventMap: Map<string, TaskState>;
  private readonly dispatch: Function;
  private readonly addLog: Function;
  private _log: HvigorLogger = HvigorLogger.getLogger(WorkerManagerImpl.name);
  private readonly maxCoreSize?: number;
  private readonly residentWorkers: Set<number>;
  private readonly cacheCapacity?: number;
  private readonly cacheTtl?: number;
  private readonly hvigorMemoryThreshold = coreParameter.properties.hvigorMemoryThreshold;

  constructor(poolConfig: PoolConfig, dispatch: Function, addLog: Function) {
    this.maxPoolNum = Math.floor(poolConfig.maxPoolNum && poolConfig.maxPoolNum >= 1 ? poolConfig.maxPoolNum : PoolConstant.MAX_POOL_NUM);
    this.minPoolNum = Math.floor(
      poolConfig.minPoolNum && poolConfig.minPoolNum >= 0 ? Math.max(poolConfig.minPoolNum, PoolConstant.MIN_POOL_NUM) : PoolConstant.MIN_POOL_NUM,
    );
    this.minPoolNum = Math.min(this.minPoolNum, this.maxPoolNum);
    this.recycleInterval = Math.floor(
      poolConfig.recycleInterval && poolConfig.recycleInterval > 0 ? poolConfig.recycleInterval : PoolConstant.DEFAULT_RECYCLE_INTERVAL,
    );
    this.maxIdleTime = Math.floor(poolConfig.maxIdleTime && poolConfig.maxIdleTime > 0 ? poolConfig.maxIdleTime : PoolConstant.DEFAULT_MAX_IDLE_TIME);
    this.idleWorkers = new Map();
    this.busyWorkers = new Map();
    this.deadWorkers = new Set();
    this.timer = setInterval(() => this.recycle(), this.recycleInterval);
    this.callbacks = new Map();
    this.eventMap = new Map([
      [PoolConstant.WORK_DONE, TaskState.END],
      [PoolConstant.WORK_ERROR, TaskState.ERROR],
    ]);
    this.dispatch = dispatch;
    this.addLog = addLog;
    this.maxCoreSize = poolConfig.maxCoreSize !== undefined && poolConfig.maxCoreSize >= 0 ? poolConfig.maxCoreSize : undefined;
    if (this.maxCoreSize) {
      this.maxCoreSize = Math.floor(this.maxCoreSize);
    }
    this.residentWorkers = new Set<number>();
    this.cacheCapacity = poolConfig.cacheCapacity;
    this.cacheTtl = poolConfig.cacheTtl;
    for (let i = 0; i < this.minPoolNum; i++) {
      this.createWorker();
    }
  }

  getMaxPoolNum(): number {
    return this.maxPoolNum;
  }

  getMinPoolNum(): number {
    return this.minPoolNum;
  }

  cleanUp(callback: () => void) {
    const _cleanUp = (workers: Map<number, HWorker>) => {
      workers.forEach((hWorker, id) => {
        // 如果处于busy状态，说明是在别的worker，或者主线程触发了失败，但是也强制此worker退出，fail fast
        if (this.residentWorkers.has(id) && !this.busyWorkers.has(id) && hWorker.hasCache()) {
          this.cleanHvigorConfig(hWorker);
          this.cleanHvigorLogger(hWorker);
          hWorker.sendMessageToWorker({ event: PoolConstant.UNMOUNT_TSC_COMMON_CACHE_EVENT });
          return;
        }
        this._log.debug(`Cleanup worker ${id}.`);
        workers.delete(id);
        hWorker.terminate(() => {
          this._log.debug(`Worker ${id} has been cleaned up.`);
          this._log.debug(`Current idle worker size: ${this.idleWorkers.size}.`);
          this._log.debug(`Current resident worker size: ${this.residentWorkers.size}.`);
          if (this.idleWorkers.size !== this.residentWorkers.size) {
            return;
          }
          callback();
        });
      });
    };

    if (this.busyWorkers.size + this.idleWorkers.size === 0) {
      this._log.debug("There's no busy workers and idle workers need cleanup.");
      this.callbacks.clear();
      callback();
    } else {
      _cleanUp(this.busyWorkers);
      _cleanUp(this.idleWorkers);
      this.callbacks.clear();
    }
  }

  clear(callback: Function): void {
    clearInterval(this.timer);
    if (this.busyWorkers.size + this.idleWorkers.size === 0) {
      this._log.debug("There's no busy workers and idle workers need to be cleared.");
      this.deadWorkers.clear();
      this.residentWorkers.clear();
      this.callbacks.clear();
      callback();
    } else {
      this.clearWorkers(callback, this.idleWorkers);
      this.clearWorkers(callback, this.busyWorkers);
      this.callbacks.clear();
    }
  }

  /**
   * 构建结束时，每个线程HVIGOR_CONFIG_LOADER进行清理
   */
  private cleanHvigorConfig(hWorker: HWorker) {
    hWorker.sendMessageToWorker({ event: PoolConstant.CLEAN_HVIGOR_CONFIG_LOADER });
  }


  /**
   * 构建结束时，每个线程HVIGOR_LOGGER进行清理
   */
  private cleanHvigorLogger(hWorker: HWorker) {
    hWorker.sendMessageToWorker({ event: PoolConstant.CLEAN_HVIGOR_LOGGER });
  }

  /**
   * 清空worker的map集合
   *
   * @param callback 清空后的回调函数
   * @param workers 待清空的worker的map集合
   */
  private clearWorkers(callback: Function, workers: Map<number, HWorker>): void {
    workers.forEach((hWorker, id) => {
      this._log.debug(`Clear worker ${id}.`);
      workers.delete(id);
      hWorker.terminate(() => {
        this._log.debug(`Worker ${id} has been cleared.`);
        this._log.debug(`Current idle worker size: ${this.idleWorkers.size}.`);
        this._log.debug(`Current busy worker size: ${this.busyWorkers.size}.`);
        if (!this.idleWorkers.size && !this.busyWorkers.size) {
          this.deadWorkers.clear();
          this.residentWorkers.clear();
          callback();
        }
      });
    });
  }

  createWorker(workerId?: number, preludeDeps: string[] = []): boolean {
    if (this.isFull()) {
      this._log.debug('Failed to create worker since the worker pool is full.');
      return false;
    }
    const id: number = this.getAllocatedId(workerId);
    if (id < 0) {
      this._log.debug('Failed to create worker since no worker id could be allocated.');
      return false;
    }
    const shouldResident = this.maxCoreSize === undefined || this.residentWorkers.size < this.maxCoreSize;
    this._log.debug(`Create ${shouldResident ? '' : 'non-'} resident worker with id: ${id}.`);
    const worker: Worker = this._createWorker(preludeDeps, shouldResident);
    const hWorker = new HWorker(worker, id);
    this.idleWorkers.set(hWorker.getId(), hWorker);
    shouldResident && this.residentWorkers.add(id);
    this.listenWorkerEvents(worker, hWorker, id);
    worker.stdout.on('data', this.getDataHandler(hWorker)(process.stdout));
    worker.stderr.on('data', this.getDataHandler(hWorker)(process.stderr));
    return true;
  }

  private listenWorkerEvents(worker: Worker, hWorker: HWorker, id: number) {
    this.listenExitEvent(worker, hWorker, id);
    this.listenErrorEvent(worker, hWorker, id);
    this.listenMessageEvent(worker, hWorker, id);
  }

  private listenExitEvent(worker: Worker, hWorker: HWorker, id: number) {
    worker.on('exit', (exitCode) => {
      this.idleWorkers.delete(id);
      this.busyWorkers.delete(id);
      this.deadWorkers.add(id);
      this.residentWorkers.delete(id);
      this._log.debug(`worker[${id}] exits with exit code ${exitCode}.`);
      this.dispatch(this, hWorker.getId());
    });
  }

  private listenErrorEvent(worker: Worker, hWorker: HWorker, id: number) {
    worker.on('error', (error) => {
      if ((error as Error & { dece: string }).dece) {
        this.errorCallback((error as Error & { dece: string }).dece);
      }
      this._log.debug(`worker[${id}] encounters error: ${error.message}.`);
      const curWorkId = this.onWorkError(hWorker);
      this.notifyError(curWorkId, PoolConstant.WORK_ERROR, error);
    });
  }

  private listenMessageEvent(worker: Worker, hWorker: HWorker, id: number) {
    worker.on('message', async (feedback) => {
      hWorker.updateCacheSize(feedback.cacheSize ?? 0);
      if (feedback.event === PoolConstant.WORK_ERROR) {
        this._log.debug(`worker[${id}] has one work error.`);
        const curWorkId = this.onWorkError(hWorker);
        const errorInfo = feedback.returnVal;
        const error = new ExtendedErrorInfo(errorInfo.message, errorInfo.extraInfo);
        if (errorInfo.stack) {
          error.stack = errorInfo.stack;
        }
        this.notifyError(curWorkId, PoolConstant.WORK_ERROR, error);
      } else {
        this._log.debug(`worker[${id}] has one work done.`);
        const tcb = TcbStore.getTCB(feedback.id) as ITCB;
        if (this.callbacks.has(feedback.id) && feedback.event === PoolConstant.WORK_DONE) {
          const callback = this.callbacks.get(feedback.id);
          try {
            await callback?.fn(...(tcb?.useReturnVal() ? [feedback.returnVal] : callback.input), id);
          } catch (e: any) {
            this.notifyError(this.onWorkError(hWorker), PoolConstant.CALLBACK_ERROR, e);
          }
        }
        hWorker.setLastBusyTime(Date.now());
        hWorker.cancelCurWorkId();
        if (feedback.startTime) {
          tcb?.setStartTime(feedback.startTime);
        }
        tcb?.setEndTime(feedback.endTime ?? Number(process.hrtime.bigint()));
        tcb?.setState(this.eventMap.get(feedback.event)!);
        this.busyWorkers.delete(hWorker.getId());
        hWorker.setState(WorkerState.IDLE);
        this.idleWorkers.set(hWorker.getId(), hWorker);
        this.notify(feedback.id, feedback.event);

        // 配置内存阈值时，任务使用worker循环分发，无配置时，分发指定线程
        if (this.hvigorMemoryThreshold && this.hvigorMemoryThreshold > 0) {
          this.dispatch(this);
        } else {
          this.dispatch(this, hWorker.getId());
        }
      }
    });
  }

  private onWorkError(hWorker: HWorker) {
    const curWorkId = hWorker.getCurWorkId()!;
    hWorker.setLastBusyTime(Date.now());
    hWorker.cancelCurWorkId();
    (TcbStore.getTCB(curWorkId) as ITCB)?.setEndTime(Number(process.hrtime.bigint()));
    (TcbStore.getTCB(curWorkId) as ITCB)?.setState(this.eventMap.get(PoolConstant.WORK_ERROR)!);
    this.busyWorkers.delete(hWorker.getId());
    hWorker.setState(WorkerState.IDLE);
    this.idleWorkers.set(hWorker.getId(), hWorker);
    return curWorkId;
  }

  private _createWorker(preludeDeps: string[] = [], shouldResident: boolean) {
    return new Worker(path.resolve(__dirname, PoolConstant.WORKER_ACTION_PATH), {
      workerData: {
        logLevel: (this._log.getLevel() as Level).levelStr,
        preludeDeps: preludeDeps,
        shouldResident: shouldResident,
        cacheCapacity: this.cacheCapacity,
        cacheTtl: this.cacheTtl,
        isDaemon: coreParameter.startParams.daemon,
      },
      env: process.env,
      stdout: true,
      stderr: true,
    });
  }

  /**
   * 获取创建worker时分配的id
   *
   * @param expectedId 希望被分配的id
   * @returns {number} 被分配的id
   */
  private getAllocatedId(expectedId?: number): number {
    let id = -1;
    if (expectedId !== undefined && (this.isBusyWorker(expectedId) || this.isIdleWorker(expectedId))) {
      return id;
    } else if (expectedId !== undefined) {
      id = expectedId;
      if (this.isRecycledWorker(expectedId)) {
        this.deadWorkers.delete(id);
      }
    } else if (this.deadWorkers.size) {
      id = this.deadWorkers.values().next().value;
      this.deadWorkers.delete(id);
    } else {
      for (let i = 0; i < this.maxPoolNum; i++) {
        if (!this.isBusyWorker(i) && !this.isIdleWorker(i)) {
          id = i;
          break;
        }
      }
    }
    return id;
  }

  /**
   * 通知执行引擎工作单元执行完毕的事件
   *
   * @param id 工作单元id
   * @param event 待通知的事件
   */
  private notify(id: string, event: string): void {
    const tcb: ITCB = TcbStore.getTCB(id) as ITCB;
    if (tcb && !tcb.isWarmUp()) {
      tcb.emit(event, id);
    }
  }

  private notifyError(id: string, event: string, error: Error) {
    const tcb: ITCB = TcbStore.getTCB(id) as ITCB;
    if (tcb && !tcb.isWarmUp()) {
      tcb.emit(event, {
        id: id,
        error: error,
      });
    }
  }

  /**
   * 获取日志数据处理函数
   *
   * @param hWorker 需要处理的日志worker对象
   * @returns {Function} 日志数据处理函数
   */
  private getDataHandler(hWorker: HWorker): Function {
    return (writeStream: (NodeJS.WriteStream & { fd: 1 }) | (NodeJS.WriteStream & { fd: 2 })) => {
      return (data: Buffer) => {
        if (!hWorker.getCurWorkId()) {
          writeStream.write(data);
          return;
        }
        const dataStr: string = iconv.decode(data, PoolConstant.ENCODING);
        const log: WorkLog = {
          type: LogType.WORK,
          time: isHvigorLogWithTime(dataStr) ? dataStr.substring(0, PoolConstant.TIME_PREFIX_LENGTH) : PoolConstant.UNKNOWN_LOG_TIME,
          workerId: hWorker.getId(),
          content: dataStr,
          taskPath: TcbStore.getTCB(hWorker.getCurWorkId()!)?.getTaskPath(),
          taskName: TcbStore.getTCB(hWorker.getCurWorkId()!)?.getTaskName(),
          taskCompletePath: TcbStore.getTCB(hWorker.getCurWorkId()!)?.getTaskCompletePath(),
        };
        this.addLog(log);
        writeStream.write(data);
      };
    };
  }
  doWork(work: Work, callback: Function, callbackInput: unknown[], workerId: number): boolean {
    if (!this.idleWorkers.has(workerId)) {
      return false;
    }
    if (
      this.busyWorkers.size !== 0 &&
      this.hvigorMemoryThreshold &&
      this.hvigorMemoryThreshold > 0 &&
      process.memoryUsage().rss / 1024 / 1024 >= this.hvigorMemoryThreshold &&
      work.isMemorySensitive()
    ) {
      this._log.debug(`A work dispatched to worker[${TcbStore.getTCB(work.getId())?.getTaskCompletePath()}] failed 
      because of process.memoryUsage().rss:${process.memoryUsage().rss / 1024 / 1024} > ${this.hvigorMemoryThreshold}.`);
      return false;
    }
    const worker: HWorker = this.idleWorkers.get(workerId)!;
    if (worker.isMeltdown() || !this.idleWorkers.get(workerId)!.doWork(work)) {
      return false;
    }
    this.idleWorkers.delete(workerId);
    this.busyWorkers.set(workerId, worker);
    if (work.hasSideEffects()) {
      worker.meltdown();
    }
    (TcbStore.getTCB(work.getId()) as ITCB).setState(TaskState.RUNNING);
    (TcbStore.getTCB(work.getId()) as ITCB).setStartTime(Number(process.hrtime.bigint()));
    (TcbStore.getTCB(work.getId()) as ITCB).setWorkerId(workerId);
    this.callbacks.set(work.getId(), { fn: callback, input: callbackInput });
    return true;
  }

  getAvailableWorkers(): number[] {
    return Array.from(this.idleWorkers.keys());
  }

  hasAvailableWorkers(): boolean {
    return this.idleWorkers.size > 0;
  }

  recycle(): void {
    this.idleWorkers.forEach((value, key) => {
      if (
        this.idleWorkers.size + this.busyWorkers.size > this.minPoolNum &&
        Date.now() - value.getLastBusyTime() > this.maxIdleTime &&
        !this.residentWorkers.has(key)
      ) {
        this.idleWorkers.delete(key);
        this.deadWorkers.add(key);
        value.terminate(noop);
      }
    });
  }

  setMaxIdleTime(maxIdleTime: number): void {
    if (maxIdleTime > 0) {
      this.maxIdleTime = maxIdleTime;
    }
  }

  setRecycleInterval(recycleInterval: number): void {
    if (recycleInterval > 0) {
      this.recycleInterval = recycleInterval;
      clearInterval(this.timer);
      this.timer = setInterval(() => {
        this.recycle();
      }, this.recycleInterval);
    }
  }

  setErrorCallback(errorCallback: Function) {
    if (this.errorCallback === noop) {
      this.errorCallback = errorCallback;
    }
  }

  isBusyWorker(workerId: number): boolean {
    return this.busyWorkers.has(workerId);
  }

  isIdleWorker(workerId: number): boolean {
    return this.idleWorkers.has(workerId);
  }

  isRecycledWorker(workerId: number): boolean {
    return this.deadWorkers.has(workerId);
  }

  isFull(): boolean {
    return this.idleWorkers.size + this.busyWorkers.size >= this.maxPoolNum;
  }
}
